# All content Copyright 2010 Cyrus Omar <cyrus.omar@gmail.com> unless otherwise
# specified.
#
# Contributors:
#     Cyrus Omar <cyrus.omar@gmail.com>
#
# This file is part of, and licensed under the terms of, the atomic-hedgehog
# package.
#
# The atomic-hedgehog package is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of the License,
# or (at your option) any later version.
#
# The atomic-hedgehog package is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
# License for more details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with the atomic-hedgehog package. If not, see <http://www.gnu.org/licenses/>.
"""Connectivity and communication stuff lives here."""

from ahh.sim import Node, ModelNode, param, Allocation, Atomics
from ahh.util import JaggedMatrix
from ahh.util.opencl import cl_int, PointerType

class AtomicSend(ModelNode):
    def __init__(self, parent, basename="AtomicSend"):
        ModelNode.__init__(self, parent, basename)
        Atomics(self.sim)

    @param
    def neighbor_size(default_hook="_pre_spike_propagation"):
        """Expression used to calculate the number of neighbors for this neuron."""

    @param
    def i_stride(default=1, cl_type=cl_int,
                 default_hook="_pre_spike_propagation"):
        """How many elements to stride when looping over neighbors."""

    @param
    def neighbors(default_hook="_pre_spike_send"):
        """The jagged matrix of neighbor data."""

    @param
    def target(cl_type=PointerType(cl_int), default_hook="_pre_spike_send"):
        """Expression used to calculate the target buffer for the spike."""

    @param
    def int_weight(default=1, cl_type=cl_int, default_hook="_pre_spike_send"):
        """Expression used to calculate the integer-valued weight for the spike."""

    def _spike_propagation(self, g):
        g.line("""for (int i=0; i < neighbor_size; i += i_stride) {""")
        g.indent_depth += 1
        self._staged_generation_hook("spike_send", g)
        g.indent_depth -= 1
        g.line("""}""")

    def _spike_send(self, g):
        g.line("atom_add(target + realization_start_idx_div + neighbors[i], int_weight);")

class AtomicReceiver(ModelNode):
    def __init__(self, parent, synapse, buffer, basename="AtomicReceive"):
        ModelNode.__init__(self, parent, basename)
        self.synapse = synapse
        self.buffer = buffer

    _default_hook = "_read_incoming_spikes"

    @param
    def weight(default=1.0, default_hook=_default_hook):
        """The weight to use."""

    @param
    def weight_transfer(default="weight*readout", default_hook=_default_hook):
        """The weight transfer expression to use."""

    @param
    def readout(default="incoming[idx]", default_hook=_default_hook):
        """The readout expression."""

    @param
    def incoming(default_hook=_default_hook):
        """The incoming buffer expression. If not set, defaults to buffer.incoming."""

    @param
    def idx(default="idx_state", default_hook=_default_hook):
        """The index expression to use for the buffers.

        If you need different ones for different buffers, add multiple units.
        """

    def _pre_finalize(self):
        if self.incoming is None:
            self.incoming = self.buffer.incoming
        self.synapse.lookup += self.weight_transfer

class AtomicBuffer(Node):
    def __init__(self, parent, length=0, basename="AtomicBuffer"):
        Node.__init__(self, parent, basename)
        self.length = length

    def _finalize(self):
        length = self.length
        buffer_even = Allocation(self, "buffer_even", cl_int, length, False,
                                     "AtomicBuffers")
        buffer_odd = Allocation(self, "buffer_odd", cl_int, length, False,
                                    "AtomicBuffers")

    def _specify_step_arguments(self, step_arguments):
        step_arguments += ("__global const int * const %s" % self.incoming,
                           "__global const int * const %s" % self.outgoing)

    @property
    def incoming(self):
        return self.name + "_in"

    @property
    def outgoing(self):
        return self.name + "_out"

    @property
    def expression(self):
        return self.name

    @classmethod
    def for_synapse(cls, synapse, parent=None, size=None):
        if parent is None:
            parent = synapse.sim

        if size is None:
            if hasattr(parent, 'count'):
                size = parent.count
            elif hasattr(parent, 'n_neurons_per_division_max'):
                size = parent.n_neurons_per_division_max

        return cls(parent, size, synapse.basename + "_transfer")
